import os
import os.path as osp
import sys
import glfw
import glm
from OpenGL.GL import *
import numpy as np
import trimesh
import ctypes
import math
import copy
from PIL import Image
from learnopengl.ibl.shader import Shader
from scipy.spatial.transform import Rotation as R
from tqdm import tqdm
from multiprocessing import Pool, cpu_count


print_device = True
SCR_WIDTH = SCR_HEIGHT = 512
crop_scale = 1.05

euler_angle = [90, 0, 0]
r = R.from_euler('xyz', euler_angle, degrees=True)
rotation_matrix = r.as_matrix()
rotation_matrix_inv = np.linalg.inv(rotation_matrix)


# Vertex Shader Source
verticesShaderSource = '''
    #version 330 core
    layout (location = 0) in vec3 aPos;
    layout (location = 1) in vec2 aTexCoord;
    layout (location = 2) in vec3 aNormal;
    
    out vec2 TexCoord;
    out vec3 Normal;
    
    uniform mat4 model;
    uniform mat4 view;
    uniform mat4 projection;
    
    void main()
    {
        Normal = mat3(transpose(inverse(model))) * aNormal;
        gl_Position = projection * view * model * vec4(aPos, 1.0);
        TexCoord = aTexCoord;
    }
'''

# Fragment Shader Source
fragmentShaderSource = '''
    #version 330 core
    out vec4 FragColor;
    
    in vec2 TexCoord;

    uniform sampler2D ourTexture;
    
    void main()
    {
        FragColor = texture(ourTexture, TexCoord);
    }
'''

# Normal_Vert Shader Source
NormalVerticesShaderSource = '''
    #version 330 core
    layout (location = 0) in vec3 aPos;
    layout (location = 1) in vec3 aNormal;
    
    out vec3 Normal;
    
    uniform mat4 model;
    uniform mat4 view;
    uniform mat4 projection;
    
    void main()
    {
        Normal = mat3(transpose(inverse(model))) * aNormal;
        gl_Position = projection * view * model * vec4(aPos, 1.0);
    }
'''

# Normal Shader Source
normalShaderSource = '''
    #version 330 core
    out vec4 FragColor;
    
    in vec3 Normal;

    void main()
    {
        vec3 normals = normalize(Normal);
        vec3 color = normals * 0.5 + 0.5;
        FragColor = vec4(color, 1.0);
    }
'''

def load_texture(texture_path, texture_unit, flip=True):
    img = Image.open(texture_path).convert('RGB')
    if flip:
        img = img.transpose(Image.FLIP_TOP_BOTTOM)
    img_data = img.tobytes()
    width, height = img.size

    texture = glGenTextures(1)
    glActiveTexture(GL_TEXTURE0 + texture_unit)
    glPixelStorei(GL_UNPACK_ALIGNMENT, 1)
    glBindTexture(GL_TEXTURE_2D, texture)

    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, GL_RGB, GL_UNSIGNED_BYTE, img_data)
    glGenerateMipmap(GL_TEXTURE_2D)
    
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAX_LEVEL, 3)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR)

    return texture

def framebuffer_size_callback(window, width, height):
    glViewport(0, 0, width, height)

def processInput(window):
    if glfw.get_key(window, glfw.KEY_ESCAPE) == glfw.PRESS:
        glfw.set_window_should_close(window, True)

def save_normal_map(width, height, filepath, type_view):
    glReadBuffer(GL_COLOR_ATTACHMENT0)
    data = glReadPixels(0, 0, width, height, GL_RGBA, GL_UNSIGNED_BYTE)
    image = Image.frombytes("RGBA", (width, height), data)
    image = image.transpose(Image.FLIP_TOP_BOTTOM)
    if type_view == 'back':
        image = image.transpose(Image.FLIP_LEFT_RIGHT)
    image.save(filepath)

def save_render_map(width, height, filepath, type_view):
    glReadBuffer(GL_COLOR_ATTACHMENT0)
    data = glReadPixels(0, 0, width, height, GL_RGBA, GL_UNSIGNED_BYTE)
    image = Image.frombytes("RGBA", (width, height), data)
    image = image.transpose(Image.FLIP_TOP_BOTTOM)
    if type_view == 'back':
        image = image.transpose(Image.FLIP_LEFT_RIGHT)
    image.save(filepath)

# 将 3x3 旋转矩阵转换为 4x4 齐次矩阵
def to_homogeneous(matrix):
    assert matrix.shape == (3, 3)
    homogeneous_matrix = np.eye(4)
    homogeneous_matrix[:3, :3] = matrix
    return homogeneous_matrix

def normalize_vector(vertices, refer_vertices=None):
    if refer_vertices is not None:
        vertices_base = refer_vertices
    else:
        vertices_base = vertices
    min_vals = vertices_base.min(axis=0)
    max_vals = vertices_base.max(axis=0)
    center = (max_vals + min_vals) / 2.0
    scale = (max_vals - min_vals).max() / 2.0

    vector_new = (vertices - center) / scale
    return vector_new


def process_clothed_files(mesh_file, smpl_file, texture_path, save_dir, SCR_WIDTH, SCR_HEIGHT, is_XRotate):
    # 加载网格和纹理
    mesh = trimesh.load(mesh_file, process=False)
    vertices_ori = np.array(mesh.vertices, dtype=np.float32)
    clothed_vertices = copy.deepcopy(vertices_ori)

    faces = np.array(mesh.faces.flatten(), dtype=np.int32)
    normals = np.array(mesh.vertex_normals, dtype=np.float32)
    uv_coords = np.array(mesh.visual.uv, dtype=np.float32)
    vertices = np.concatenate([vertices_ori, uv_coords, normals], axis=1)

    # 计算正交投影的坐标
    # if is_XRotate:
    #     clothed_vertices = np.matmul(rotation_matrix_inv, clothed_vertices.T).T
    #     smpl_mesh = trimesh.load(smpl_file, process=False)
    #     vertices_smpl = np.array(smpl_mesh.vertices, dtype=np.float32)
    #     vertices_smpl = np.matmul(rotation_matrix_inv, vertices_smpl.T).T
    #     vertices_new = normalize_vector(vertices_smpl)
    # else:
    #     # vertices_new = vertices_ori
    #     vertices_new = clothed_vertices
    
    vertices_new = clothed_vertices
    min_x = vertices_new[:, 0].min()
    max_x = vertices_new[:, 0].max()
    min_y = vertices_new[:, 1].min()
    max_y = vertices_new[:, 1].max()
    
    # left, right, bottom, top
    w = max_x - min_x
    h = max_y - min_y
    mean_w = (max_x + min_x) / 2
    mean_h = (max_y + min_y) / 2
    if w > h:
        # w = np.clip(w * crop_scale, -1.0, 1.0)
        w = w * crop_scale
        left = mean_w - w / 2
        right = mean_w + w / 2
        top = mean_h + w / 2
        bottom = mean_h - w / 2
    else:
        # h = np.clip(h * crop_scale, -1.0, 1.0)
        h = h * crop_scale
        left = mean_w - h / 2
        right = mean_w + h / 2
        top = mean_h + h / 2
        bottom = mean_h - h / 2
    zNear = -100
    zFar = 100.0

    # 初始化
    if not glfw.init():
        raise Exception("GLFW can't be initialized")
    glfw.window_hint(glfw.VISIBLE, glfw.FALSE)  # 隐藏窗口
    glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 3)
    glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
    glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
    
    window = glfw.create_window(SCR_WIDTH, SCR_HEIGHT, "LearnOpenGL", None, None)
    if not window:
        glfw.terminate()
        raise Exception("GLFW window can't be created")

    glfw.make_context_current(window)
    glfw.set_framebuffer_size_callback(window, framebuffer_size_callback)
    
    glfw.set_input_mode(window, glfw.CURSOR, glfw.CURSOR_DISABLED)
    
    glEnable(GL_DEPTH_TEST)
    
    # 打印 GPU 信息
    global print_device
    if print_device:
        print("Renderer:", glGetString(GL_RENDERER).decode())
        print("OpenGL version supported:", glGetString(GL_VERSION).decode())
        print_device = False
    
    VAO = glGenVertexArrays(1)
    VBO = glGenBuffers(1)
    EBO = glGenBuffers(1)
    
    glBindVertexArray(VAO)
    
    glBindBuffer(GL_ARRAY_BUFFER, VBO)
    glBufferData(GL_ARRAY_BUFFER, vertices.nbytes, vertices, GL_STATIC_DRAW)
    
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO)
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, faces.nbytes, faces, GL_STATIC_DRAW)

    stride = vertices.itemsize * 8  # 3 for pos, 2 for tex, 3 for normal
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, stride, ctypes.c_void_p(0))
    glEnableVertexAttribArray(0)
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, stride, ctypes.c_void_p(3 * vertices.itemsize))
    glEnableVertexAttribArray(1)
    glVertexAttribPointer(2, 3, GL_FLOAT, GL_FALSE, stride, ctypes.c_void_p(5 * vertices.itemsize))
    glEnableVertexAttribArray(2)
    
    renderShaderProgram = Shader(verticesShaderSource, fragmentShaderSource)
    normalShaderProgram = Shader(verticesShaderSource, normalShaderSource)
    
    texture_map = load_texture(texture_path, 0)
    
    renderShaderProgram.use()
    glUniform1i(glGetUniformLocation(renderShaderProgram.ID, "ourTexture"), 0)

    # 设置帧缓冲对象
    FBO = glGenFramebuffers(1)
    glBindFramebuffer(GL_FRAMEBUFFER, FBO)
    
    # 创建纹理
    textureColorBuffer = glGenTextures(2)
    glBindTexture(GL_TEXTURE_2D, textureColorBuffer[0])
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, SCR_WIDTH, SCR_HEIGHT, 0, GL_RGBA, GL_UNSIGNED_BYTE, None)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, textureColorBuffer[0], 0)
    
    glBindTexture(GL_TEXTURE_2D, textureColorBuffer[1])
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, SCR_WIDTH, SCR_HEIGHT, 0, GL_RGBA, GL_UNSIGNED_BYTE, None)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT1, GL_TEXTURE_2D, textureColorBuffer[1], 0)
    
    # 创建渲染缓冲对象
    rbo = glGenRenderbuffers(1)
    glBindRenderbuffer(GL_RENDERBUFFER, rbo)
    glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH24_STENCIL8, SCR_WIDTH, SCR_HEIGHT)
    glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_STENCIL_ATTACHMENT, GL_RENDERBUFFER, rbo)
    
    if glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE:
        raise Exception("Framebuffer is not complete!")
    
    glBindFramebuffer(GL_FRAMEBUFFER, 0)
    
    y_rotate = np.arange(0, 360, 360 // 36)
    i = 0
    while not glfw.window_should_close(window):  
        processInput(window)
        
        model = glm.rotate(glm.mat4(1.0), math.radians(y_rotate[i]), glm.vec3(0, 1, 0))
        if is_XRotate:
            model = model @ glm.mat4(to_homogeneous(rotation_matrix_inv))
        projection = glm.ortho(left, right, bottom, top, zNear, zFar)
        view_front = glm.lookAt(glm.vec3(0, 0, 1), glm.vec3(0, 0, 0), glm.vec3(0, 1, 0))
        view_back = glm.lookAt(glm.vec3(0, 0, -1), glm.vec3(0, 0, 0), glm.vec3(0, 1, 0))

        def process_front_back(view, type_view="front"):
            glBindFramebuffer(GL_FRAMEBUFFER, FBO)
            glClearColor(0.0, 0.0, 0.0, 0.0)

            # # 渲染法线贴图
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            normalShaderProgram.use()
            normalShaderProgram.set_mat4("projection", projection)
            normalShaderProgram.set_mat4("view", view)
            normalShaderProgram.set_mat4("model", model)
            
            glBindVertexArray(VAO)
            glDrawElements(GL_TRIANGLES, len(faces), GL_UNSIGNED_INT, None)
            os.makedirs(f'{save_dir}/view_{i:02d}', exist_ok=True)
            save_normal_map(SCR_WIDTH, SCR_HEIGHT, f'{save_dir}/view_{i:02d}/normal_{type_view}.png', type_view)

            # 渲染纹理贴图
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            renderShaderProgram.use()
            renderShaderProgram.set_mat4("projection", projection)
            renderShaderProgram.set_mat4("view", view)
            renderShaderProgram.set_mat4("model", model)
            
            glActiveTexture(GL_TEXTURE0)
            glBindTexture(GL_TEXTURE_2D, texture_map)
            
            glBindVertexArray(VAO)
            glDrawElements(GL_TRIANGLES, len(faces), GL_UNSIGNED_INT, None)
            save_render_map(SCR_WIDTH, SCR_HEIGHT, f'{save_dir}/view_{i:02d}/image_{type_view}.png', type_view)

            glBindFramebuffer(GL_FRAMEBUFFER, 0)
        
        process_front_back(view_front, "front")
        process_front_back(view_back, "back")

        glfw.swap_buffers(window)
        glfw.poll_events()
        i += 1
        if i == len(y_rotate):
        # if i == 1:
            break
    
    glDeleteBuffers(1, [VBO])
    glDeleteBuffers(1, [EBO])
    glDeleteVertexArrays(1, [VAO])
    glDeleteTextures(1, textureColorBuffer)
    glDeleteRenderbuffers(1, [rbo])
    glDeleteFramebuffers(1, [FBO])
    glfw.terminate()


def process_smpl_files(mesh_file, smpl_file, save_dir, SCR_WIDTH, SCR_HEIGHT, is_XRotate):
    # 加载网格和纹理
    mesh = trimesh.load(smpl_file, process=False)
    vertices_ori = np.array(mesh.vertices, dtype=np.float32)

    clothed_mesh = trimesh.load(mesh_file, process=False)
    clothed_vertices = np.array(clothed_mesh.vertices, dtype=np.float32)

    faces = np.array(mesh.faces.flatten(), dtype=np.uint32)
    normals = np.array(mesh.vertex_normals, dtype=np.float32)
    vertices = np.concatenate([vertices_ori, normals], axis=1)
    
    # 计算正交投影的坐标
    # if is_XRotate:
        # clothed_vertices = np.matmul(rotation_matrix_inv, clothed_vertices.T).T
        # smpl_mesh = copy.deepcopy(mesh)
        # vertices_smpl = np.array(smpl_mesh.vertices, dtype=np.float32)
        # vertices_new = normalize_vector(vertices_smpl)
    # else:
    #     # vertices_new = vertices_ori
    #     vertices_new = clothed_vertices

    vertices_new = clothed_vertices
    min_x = vertices_new[:, 0].min()
    max_x = vertices_new[:, 0].max()
    min_y = vertices_new[:, 1].min()
    max_y = vertices_new[:, 1].max()
    
    # left, right, bottom, top
    w = max_x - min_x
    h = max_y - min_y
    mean_w = (max_x + min_x) / 2
    mean_h = (max_y + min_y) / 2
    if w > h:
        # w = np.clip(w * crop_scale, -1.0, 1.0)
        w = w * crop_scale
        left = mean_w - w / 2
        right = mean_w + w / 2
        top = mean_h + w / 2
        bottom = mean_h - w / 2
    else:
        # h = np.clip(h * crop_scale, -1.0, 1.0)
        h = h * crop_scale
        left = mean_w - h / 2
        right = mean_w + h / 2
        top = mean_h + h / 2
        bottom = mean_h - h / 2
    zNear = -100
    zFar = 100.0

    # 初始化
    if not glfw.init():
        raise Exception("GLFW can't be initialized")
    glfw.window_hint(glfw.VISIBLE, glfw.FALSE)  # 隐藏窗口
    glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 3)
    glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
    glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
    
    window = glfw.create_window(SCR_WIDTH, SCR_HEIGHT, "LearnOpenGL", None, None)
    if not window:
        glfw.terminate()
        raise Exception("GLFW window can't be created")

    glfw.make_context_current(window)
    glfw.set_framebuffer_size_callback(window, framebuffer_size_callback)
    
    glfw.set_input_mode(window, glfw.CURSOR, glfw.CURSOR_DISABLED)
    
    glEnable(GL_DEPTH_TEST)
    
    # 打印 GPU 信息
    global print_device
    if print_device:
        print("Renderer:", glGetString(GL_RENDERER).decode())
        print("OpenGL version supported:", glGetString(GL_VERSION).decode())
        print_device = False
    
    VAO = glGenVertexArrays(1)
    VBO = glGenBuffers(1)
    EBO = glGenBuffers(1)
    
    glBindVertexArray(VAO)
    
    glBindBuffer(GL_ARRAY_BUFFER, VBO)
    glBufferData(GL_ARRAY_BUFFER, vertices.nbytes, vertices, GL_STATIC_DRAW)
    
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO)
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, faces.nbytes, faces, GL_STATIC_DRAW)

    stride = vertices.itemsize * 6  # 3 for pos, 3 for normal
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, stride, ctypes.c_void_p(0))
    glEnableVertexAttribArray(0)
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, stride, ctypes.c_void_p(3 * vertices.itemsize))
    glEnableVertexAttribArray(1)
    
    normalShaderProgram = Shader(NormalVerticesShaderSource, normalShaderSource)
    
    # 设置帧缓冲对象
    FBO = glGenFramebuffers(1)
    glBindFramebuffer(GL_FRAMEBUFFER, FBO)
    
    # 创建纹理
    textureColorBuffer = glGenTextures(2)
    glBindTexture(GL_TEXTURE_2D, textureColorBuffer[0])
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, SCR_WIDTH, SCR_HEIGHT, 0, GL_RGBA, GL_UNSIGNED_BYTE, None)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, textureColorBuffer[0], 0)
    
    glBindTexture(GL_TEXTURE_2D, textureColorBuffer[1])
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, SCR_WIDTH, SCR_HEIGHT, 0, GL_RGBA, GL_UNSIGNED_BYTE, None)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT1, GL_TEXTURE_2D, textureColorBuffer[1], 0)
    
    # 创建渲染缓冲对象
    rbo = glGenRenderbuffers(1)
    glBindRenderbuffer(GL_RENDERBUFFER, rbo)
    glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH24_STENCIL8, SCR_WIDTH, SCR_HEIGHT)
    glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_STENCIL_ATTACHMENT, GL_RENDERBUFFER, rbo)
    
    if glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE:
        raise Exception("Framebuffer is not complete!")
    
    glBindFramebuffer(GL_FRAMEBUFFER, 0)
    
    y_rotate = np.arange(0, 360, 360 // 36)
    i = 0
    while not glfw.window_should_close(window):  
        processInput(window)
        
        model = glm.rotate(glm.mat4(1.0), math.radians(y_rotate[i]), glm.vec3(0, 1, 0))
        if is_XRotate:
            model = model @ glm.mat4(to_homogeneous(rotation_matrix_inv))
        projection = glm.ortho(left, right, bottom, top, zNear, zFar)
        view_front = glm.lookAt(glm.vec3(0, 0, 1), glm.vec3(0, 0, 0), glm.vec3(0, 1, 0))
        view_back = glm.lookAt(glm.vec3(0, 0, -1), glm.vec3(0, 0, 0), glm.vec3(0, 1, 0))

        def process_front_back(view, type_view="front"):
            glBindFramebuffer(GL_FRAMEBUFFER, FBO)
            glClearColor(0.0, 0.0, 0.0, 0.0)

            # # 渲染法线贴图
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            normalShaderProgram.use()
            normalShaderProgram.set_mat4("projection", projection)
            normalShaderProgram.set_mat4("view", view)
            normalShaderProgram.set_mat4("model", model)
            
            glBindVertexArray(VAO)
            glDrawElements(GL_TRIANGLES, len(faces), GL_UNSIGNED_INT, None)
            os.makedirs(f'{save_dir}/view_{i:02d}', exist_ok=True)
            save_normal_map(SCR_WIDTH, SCR_HEIGHT, f'{save_dir}/view_{i:02d}/T_normal_{type_view}.png', type_view)

            glBindFramebuffer(GL_FRAMEBUFFER, 0)
        
        process_front_back(view_front, "front")
        process_front_back(view_back, "back")

        glfw.swap_buffers(window)
        glfw.poll_events()
        i += 1
        if i == len(y_rotate):
        # if i == 1:
            break
    
    glDeleteBuffers(1, [VBO])
    glDeleteBuffers(1, [EBO])
    glDeleteVertexArrays(1, [VAO])
    glDeleteTextures(1, textureColorBuffer)
    glDeleteRenderbuffers(1, [rbo])
    glDeleteFramebuffers(1, [FBO])
    glfw.terminate()


def process_clothed_task(args):
    return process_clothed_files(*args)

def process_smpl_task(args):
    return process_smpl_files(*args)

def process_task(args):
    process_smpl_files(*args[1])
    # process_clothed_files(*args[0])

if __name__ == "__main__":
    start = 532
    rotate_i = 526
    src_dir = r"THuman2.1_Release"
    save_dir = r"save_images4"
    # save_dir = r"save_images1"
    input_dir = osp.join(src_dir, 'model')
    smpl_dir = osp.join(src_dir, 'smplx')

    filedir_list = os.listdir(input_dir)
    filedir_list.sort()

    file_groups = []
    for file_i in filedir_list:
        mesh_file = osp.join(input_dir, file_i, file_i+'.obj')
        smpl_file = osp.join(smpl_dir, file_i, 'mesh_smplx.obj')
        texture_path = osp.join(input_dir, file_i, 'material0.jpeg')
        assert osp.exists(mesh_file), mesh_file
        assert osp.exists(texture_path), texture_path
        file_groups.append({'mesh_file': mesh_file, 'smpl_file': smpl_file, 'texture_path': texture_path})

    os.makedirs(save_dir, exist_ok=True)
    exist_list = os.listdir(save_dir)
    exist_id = []
    # for img in exist_list:
    #     exist_id.append(int(img))

    clothed_tasks = []
    smpl_task = []
    task = []
    for i, group in enumerate(file_groups):
        if i < start:
            continue
        if i < rotate_i:
            is_XRotate = False
        else:
            is_XRotate = True
        if i in exist_id:
            continue
        save_sub_dir = os.path.join(save_dir, os.path.basename(group['mesh_file']).split('.')[0])
        clothed_tasks.append((group['mesh_file'], group['smpl_file'], group['texture_path'], save_sub_dir, SCR_WIDTH, SCR_HEIGHT, is_XRotate))
        smpl_task.append((group['mesh_file'], group['smpl_file'], save_sub_dir, SCR_WIDTH, SCR_HEIGHT, is_XRotate))
        task.append([clothed_tasks[-1], smpl_task[-1]])
    

    
    num_workers = cpu_count()  # 获取系统中的 CPU 核心数
    print(f"Number of workers: {num_workers}")

    process_func = process_task
    
    # with Pool(num_workers) as pool:
    #     for _ in tqdm(pool.imap_unordered(process_func, task), total=len(task)):
    #         pass
    # exit()
    
    # file_groups = file_groups[:5]
    # run
    for i, group in enumerate(task):
        process_func(group)
        exit()

